#include "CudaColorBar.cuh"
#include "CudaGetFluxMap.cuh"
#include <algorithm/rayTracing/qmcrt_util/check_cuda.h>
#include <cstdio>

using namespace solar::cuda;

__device__ static auto clamp(float x, float a, float b) -> float { return max(a, min(b, x)); }

__device__ static auto rgbToInt(float r, float g, float b) -> float4
{
    r = clamp(r, 0.0f, 255.0f);
    g = clamp(g, 0.0f, 255.0f);
    b = clamp(b, 0.0f, 255.0f);
    printf("r=%f, g=%f, b=%f\n", r, g, b);
    return make_float4(r, g, b, 125);
}

__device__ static auto getColorAtValueCuda(int3* colors, float* factors, float factor,
                                           int color_num) -> uchar4
{
    int3 color;
    int id = 0;
    while (factor > factors[id] && id < color_num)
    {
        id++;
    }
    // printf("id=%d, factor=%f\n",id, factor);

    if (id == color_num)
        color = colors[color_num - 1];
    else if (id == 0)
        color = colors[0];
    else
    {
        int3 left = colors[id - 1];
        int3 right = colors[id];
        float coef = (factor - factors[id - 1]) / (factors[id] - factors[id - 1]);
        float r = left.x + coef * (right.x - left.x);
        float g = left.y + coef * (right.y - left.y);
        float b = left.z + coef * (right.z - left.z);
        color = make_int3((int)r, (int)g, (int)b);
        // if((color.x==0 && color.y==0&& color.z==0) || color.x > 255 || color.y > 255 || color.z >
        // 255) printf("factor=%f,id=%d,left=(%f, %d,%d,%d), right=(%f, %d,%d,%d),
        // color=(%d,%d,%d)\n",
        //      factor, id,        factors[id-1], left.x,left.y,left.z,
        //                         factors[id], right.x, right.y, right.z,
        //                                     color.x, color.y, color.z);
    }

    return make_uchar4(color.x, color.y, color.z, 255);
}
/**
 * GPU中：图像rgba、flux、分辨率、colorbar（colorbar已经设置过最大最小值）
 */
static __global__ void cudaGetFluxMap(uchar4* img, float* data, int width, int height, int3* colors,
                                      float* factors, int color_num, float max_value)
{
    // printf("1\n");
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int bw = blockDim.x;
    int bh = blockDim.y;
    int x = blockIdx.x * bw + tx;
    int y = blockIdx.y * bh + ty;
    // printf("x=%d, y=%d, width=%d, height=%d\n", x,y, width, height);

    if (x >= width || y >= height)
        return;

    int idx = y * width + x;
    // printf("x=%d, y=%d, idx=%d\n",x,y,idx);
    img[idx] = getColorAtValueCuda(colors, factors, data[idx] / max_value, color_num);
    // int r = colorbar.getRedComponentAtValue(data[idx]);
    // int g = colorbar.getGreenComponentAtValue(data[idx]);
    // int b = colorbar.getBlueComponentAtValue(data[idx]);

    // img[idx] = make_float4((float)r / 255.0, (float)g / 255.0, (float)b / 255.0, 1.0);
    // if(img[idx].x<1e-2 && img[idx].y<1e-2 && img[idx].z<1e-2)
    // printf("data=%f, max_value=%f, width=%d, height=%d,img[%d]: %f %f %f\n",
    //         data[idx], max_value,width, height, idx, img[idx].x, img[idx].y, img[idx].z);
    // int width, int height
}

CudaGetFluxMap::~CudaGetFluxMap()
{
    if (colors_)
    {
        checkCudaErrors(cudaFree(colors_));
        colors_ = nullptr;
    }
    if (factors_)
    {
        checkCudaErrors(cudaFree(factors_));
        factors_ = nullptr;
    }
}

void CudaGetFluxMap::setColor()
{
    colorbar_ = new CudaColorBar();
    colorbar_->setColor();
    // colorbar_->test();

    color_num_ = colorbar_->getColorNum();
    checkCudaErrors(cudaMalloc((void**)&colors_, (color_num_) * sizeof(int3)));
    checkCudaErrors(cudaMalloc((void**)&factors_, (color_num_) * sizeof(float)));

    checkCudaErrors(cudaMemcpy((void*)colors_, (void*)colorbar_->getColors(),
                               color_num_ * sizeof(int3), cudaMemcpyHostToDevice));
    checkCudaErrors(cudaMemcpy((void*)factors_, (void*)colorbar_->getFactors(),
                               color_num_ * sizeof(float), cudaMemcpyHostToDevice));
}

void CudaGetFluxMap::calculateFluxMap(float* d_data, int2 resolution, float max_value,
                                      uchar4*& d_fluxmap_data)
{
    checkCudaErrors(
        cudaMalloc((void**)&d_fluxmap_data, resolution.x * resolution.y * sizeof(uchar4)));
    dim3 blocks(32, 32);
    dim3 grids((resolution.x + blocks.x - 1) / blocks.x, (resolution.y + blocks.y - 1) / blocks.y);
    cudaGetFluxMap<<<grids, blocks>>>(d_fluxmap_data, d_data, resolution.x, resolution.y, colors_,
                                      factors_, color_num_, max_value);
}
